from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
from utils import get_file_lines_count
from collections import namedtuple
import re
from typing import NamedTuple
from tqdm import tqdm
import os

TRACE_ON = "TRACE_ON"
PE_CONTEXT = re.compile(r"OCSD_GEN_TRC_ELEM_PE_CONTEXT\((.+)\)")
EL1N = "EL1N"
EL1S = "EL1S"
EL2N = "EL2N"
EL2S = "EL2S"
EL3N = "EL3N"
EL3S = "EL3S"
EL4N = "EL4N"
EL4S = "EL4S"
ISA_A64 = "ISA=A64"
ISA_A32 = "ISA=A32"
VMID = re.compile(r"VMID=(0x\w+);")
CTXTID = re.compile(r"CTXTID=(0x\w+);")
EXCEPTION_RET = "EXCEPTION_RET()"
ADDR_NACC = "ADDR_NACC"
TRACE_BEGIN = "init-decoder"
TRACE_DONE = re.compile(r"Trace buffer done, processed (\d+) bytes.")
INSTR_RANGE = re.compile(
    r"OCSD_GEN_TRC_ELEM_INSTR_RANGE\(exec range=0x(\w+):\[0x(\w+)\] num_i\((\d+)\) last_sz\((\d+)\) \(ISA=A64\) (.+)\)"
)
TIMESTAMP = re.compile(r"OCSD_GEN_TRC_ELEM_TIMESTAMP\( \[ TS=(\w+)\]; \)")
SYNC_MARKER = "OCSD_GEN_TRC_ELEM_SYNC_MARKER"
AARCH64_BR = {
    "b": {"b", "br"},
    "b+link": {"bl", "blr"},
    "<cond>": {
        "b.eq",
        "b.ne",
        "b.cs",
        "b.hs",
        "b.cc",
        "b.lo",
        "b.mi",
        "b.pl",
        "b.vs",
        "b.vc",
        "b.hi",
        "b.ls",
        "b.ge",
        "b.lt",
        "b.gt",
        "b.le",
        "b.le",
        "b.al",
        "b.nv",
        "cbnz",
        "cbz",
        "tbz",
        "tbnz",
    },
}
AARCH64_IBR = {
    "A64:ret": {"ret"},
    "A64:eret": {"eret"},
    "b": {"b", "br"},
    "b+link": {"bl", "blr"},
}
AARCH64_ISB = {"isb"}
AARCH64_DXB = {"dsb", "dmb"}
AARCH64_WFX = {"wfi", "wfe"}
AARCH64_TSTART = {"tstart"}
INSTR_CMD = {
    "BR": AARCH64_BR,
    "iBR": AARCH64_IBR,
    "ISB": AARCH64_ISB,
    "DSB.DMB": AARCH64_DXB,
    "WFI.WFE": AARCH64_WFX,
    "TSTART": AARCH64_TSTART,
}


class TraceCmd:
    def __init__(self, cmd: str):
        self.__cmd = cmd
        self.exec = False
        self.__type = None
        self.__sub_type = None
        self.__parse()

    def __parse(self):
        parts = self.__cmd.split(" ")
        cmds = [part for part in parts if part.strip() != ""]
        if cmds[1] not in INSTR_CMD:
            return
        self.exec = cmds[0] == "E"
        self.__type = cmds[1]
        try:
            self.__sub_type = cmds[2]
        except IndexError:
            if self.__type == "BR" or self.__type == "iBR":
                self.__sub_type = "b"
            else:
                self.__sub_type = None

    def __eq__(self, value: object) -> bool:
        cmd = value.split("\t")[0]
        if self.__sub_type:
            return cmd in INSTR_CMD[self.__type][self.__sub_type]
        return cmd in INSTR_CMD[self.__type]


ParserData = NamedTuple("ParserData", [("saddr", str), ("eaddr", str), ("cmd", str)])


class TraceDecode:
    def __init__(self, trace_file: str, batch_size: int = 20000):
        self.__trace_file = trace_file
        self.__batch_size = batch_size
        self.__timestamp = self.__have_timestamp()
        with open(self.__trace_file, "r") as f:
            self.__decode = f.readlines()
        self.__trace_downsizing()
        self.__lines_count = len(self.__decode)
        self.__batch_num = (
            self.__lines_count + self.__batch_size - 1
        ) // self.__batch_size
        self.__output: list[ParserData] = []
        self.__trace_parser()
        del self.__decode

    def __have_timestamp(self):
        with open(self.__trace_file, "r") as f:
            return TIMESTAMP.search(f.read(1024 * 1024)) is not None

    def __sort_by_timestamp(self):
        start = []
        end = []
        is_start = False
        max_timestamp = 0
        max_index = 0
        min_index = 0

        if not self.__timestamp:
            return

        for i in tqdm(range(len(self.__decode)), desc="Sort by timestamp"):
            try:
                timestamp = int(TIMESTAMP.search(self.__decode[i]).groups()[0], 16)
                if timestamp >= max_timestamp:
                    max_timestamp = timestamp
                    max_index = i
                    if is_start:
                        raise IndexError("start is already found")
                elif not is_start:
                    is_start = True
                    min_index = i

            except AttributeError:
                if SYNC_MARKER in self.__decode[i]:
                    continue
                if is_start:
                    start.append(self.__decode[i])
                else:
                    end.append(self.__decode[i])
        if is_start:
            for i in range(len(end)):
                last = end.pop()
                if last == self.__decode[max_index - 1]:
                    end.append(last)
                    break
            for i in range(len(start)):
                first = start.pop(0)
                if first == self.__decode[min_index + 1]:
                    start.insert(0, first)
                    break
            self.__decode = end[0:1] + start[0:-1] + end[1:] + [start[-1]]
        else:
            self.__decode = end

    def __trace_downsizing(self):
        start = 0
        for i in range(50):
            if TRACE_BEGIN in self.__decode[i]:
                start = i
                break

        self.__decode = self.__decode[start:]
        self.__sort_by_timestamp()

    def __PE_CONTEXT(self, data: str):
        exception_level = ""
        instruction_type = ""
        output = ""

        # check the exception level
        if data.find(EL1N) > 0:
            exception_level = "Following trace captured at Non-secure EL1. "
        elif data.find(EL1S) > 0:
            exception_level = "Following trace captured at Secure EL1. "
        elif data.find(EL2N) > 0:
            exception_level = "Following trace captured at Non-secure EL2. "
        elif data.find(EL2S) > 0:
            exception_level = "Following trace captured at Secure EL2. "
        elif data.find(EL3N) > 0:
            exception_level = "Following trace captured at Non-secure EL3. "
        elif data.find(EL3S) > 0:
            exception_level = "Following trace captured at Secure EL3. "
        elif data.find(EL4N) > 0:
            exception_level = "Following trace captured at Non-secure EL4. "
        elif data.find(EL4S) > 0:
            exception_level = "Following trace captured at Secure EL4. "

        # find type of instructions being traced
        if data.find(ISA_A64) > 0:
            instruction_type = "Tracing AArch64 instructions. "
        elif data.find(ISA_A32) > 0:
            instruction_type = "Tracing AArch32 instruction. "
        output = exception_level + instruction_type
        vmid = VMID.search(data)
        if vmid:
            vmid = vmid.groups()[0]
            output += "Virtual Context identifier is " + vmid + ". "

        ctxid = CTXTID.search(data)
        if ctxid:
            ctxid = ctxid.groups()[0]
            output += "Context identifier is " + ctxid + ". "

        return output

    def __trace_parser_handle(self, decode: list[str]):
        output: list[ParserData] = []
        for data in decode:
            tmp = INSTR_RANGE.search(data)
            try:
                s_addr, e_addr, num, lsz, lcmd = tmp.groups()
                output.append(ParserData(s_addr, e_addr, lcmd))
                continue
            except:
                if "OCSD_GEN_TRC_ELEM_INSTR_RANGE" in data:
                    raise Exception("Error parsing INSTR_RANGE trace")

            tmp = PE_CONTEXT.search(data)
            try:
                tmp = tmp.groups()[0]
                output.append(ParserData("PE_CONTEXT", None, self.__PE_CONTEXT(data)))
                continue
            except:
                if "OCSD_GEN_TRC_ELEM_PE_CONTEXT" in data:
                    raise Exception("Error parsing PE_CONTEXT trace")

            if TRACE_ON in data:
                output.append(ParserData(TRACE_ON, None, "Trace enabled"))
                continue

            # following if statement determines if an exception return has occured
            if EXCEPTION_RET in data:
                output.append(
                    ParserData("EXCEPTION_RET", None, "An exception return occured.")
                )

            # following if statement checks for trace binary corruption
            if ADDR_NACC in data:
                output.append(
                    ParserData(
                        "ADDR_NACC",
                        None,
                        "Tracing inaccessible memory area. Trace binary corrupted.",
                    )
                )

            # following if statement checks if trace has finished
            tmp = TRACE_DONE.search(data)
            try:
                byte = tmp.groups()[0]
                output.append(
                    ParserData(
                        "TRACE_FINISHED",
                        None,
                        "Trace decode has finished. Processed " + byte + " bytes",
                    )
                )
            except:
                if "Trace buffer done" in data:
                    raise Exception("Error parsing TRACE_DONE trace")
        return output

    def __trace_parser(self):
        # multi thread
        with ThreadPoolExecutor() as executor:
            pbar = tqdm(
                executor.map(
                    lambda batch: self.__trace_parser_handle(batch),
                    [
                        self.__decode[i : i + self.__batch_size]
                        for i in range(0, self.__lines_count, self.__batch_size)
                    ],
                ),
                desc="Decoding trace",
                total=self.__batch_num,
            )
            pbar.update(0)
            for decodeData in pbar:
                self.__output.extend(decodeData)

    def getTrace(self):
        return self.__output


if __name__ == "__main__":
    trace = TraceDecode("111")
    print("done")
